import torch
import torch.backends.cudnn as cudnn
import configparser as cp

from torch.utils.data import DataLoader

from lib.net.facenet.Facenet import Facenet
from lib.utils.dataloader import LFWDataset
from lib.utils.utils import getRootPath
from lib.utils.metrics import test

if __name__ == "__main__":
    # 获取项目根目录
    rootPath = getRootPath()
    # 读取配置文件
    filename = r'../config.ini'
    inifile = cp.ConfigParser()
    inifile.read(filename, 'UTF-8')
    # 设置参数值  详细说明可以查看config.ini配置文件
    Cuda = eval(inifile.get('test', 'Cuda'))
    input_shape = eval(inifile.get('test', 'input_shape'))
    backbone = inifile.get('test', 'backbone')
    model_path = rootPath + inifile.get('test', 'model_path')
    lfw_pairs_path = rootPath + inifile.get('test', 'lfw_pairs_path')
    batch_size = eval(inifile.get('test', 'batch_size'))
    log_interval = eval(inifile.get('test', 'log_interval'))
    png_save_path = rootPath + inifile.get('test', 'png_save_path')

    # ---------------------------------------#
    #   构建数据集加载器。
    # ---------------------------------------#
    test_dataset = LFWDataset(pairs_path=lfw_pairs_path, image_size=input_shape)

    # 用以LFW测试的数据
    test_loader = DataLoader(test_dataset, shuffle=True, batch_size=batch_size)

    model = Facenet(backbone=backbone)

    print('Loading weights into state dict...')
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    model.load_state_dict(torch.load(model_path, map_location=device), strict=False)
    model  = model.eval()

    if Cuda:
        model = torch.nn.DataParallel(model)
        cudnn.benchmark = True
        model = model.cuda()

    test(test_loader, model, png_save_path, log_interval, batch_size, Cuda)